import time
import os
import numpy as np
import queue

from flask import Flask, request, render_template, jsonify
from keras.models import Model
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers.pooling import GlobalAveragePooling2D
from keras.preprocessing.image import img_to_array, load_img
from keras.applications.densenet import DenseNet169
from keras.applications.densenet import preprocess_input
from keras import backend as K

# __name__的值就是当前文件所处的文件夹名
app = Flask(__name__)
# 建立一个全局消息队列，用于向前台发送数据 格式{token:queue},每次请求产生唯一token，实现并发功能
message_q_dict = {}
# 瑕疵类型数组
flaw_names = ['扎洞', '其他', '毛斑', '擦洞', '毛洞', '织稀', '吊经', '缺经', '跳花', '油/污渍', '正常']
# 模型参数
model_path = 'model_weight.10-0.7704_balance.hdf5'
# 图片大小
in_size = [550, 600]  # [H,W]


def get_model(INT_HEIGHT, IN_WIDTH, last_layer=11):
    """
    获得模型
    """
    base_model = DenseNet169(include_top=False, weights='imagenet', input_shape=(INT_HEIGHT, IN_WIDTH, 3))

    x = base_model.output
    x = Dropout(0.5)(x)
    x = GlobalAveragePooling2D()(x)
    predictions = Dense(last_layer, activation='softmax')(x)  # new softmax layer
    model = Model(input=base_model.input, output=predictions)
    model.summary()

    return model


# 初始化模型
K.set_learning_phase(0)
model = get_model(in_size[0], in_size[1])
model.load_weights(model_path)
model.predict(np.zeros((1, in_size[0], in_size[1], 3)))
print("模型加载完成")


# @app.route('xxxxxx')表示前端访问引号内部的网址,后台就会执行下面的函数内代码
@app.route('/')
def index():
    # 将index.html页面返回给前端，默认页面会在templates文件夹下寻找
    return render_template("index.html")


@app.route('/start')
def start():
    # 获取前端name为pic_dir的输入框中的值
    # request.args.get('XX')能获取到前端传过来的两种值
    # 1、是表单内name对应的值，此处是这种
    # 2、是网址末尾添加的值，例如/work?aaa=11,例如get('aaa')就能拿到11，下面一个函数是这种情况
    pic_dir = request.args.get('pic_dir')

    pic_cnt = 0
    for parent, _, files in os.walk(pic_dir):
        for file in files:
            if file[-3:] == 'jpg':
                pic_cnt += 1
    # 返回start.html页面,同时传两个值到前端，前端可以通过{{pic_dir}}和{{pic_cnt}}取到相应值
    # 同时还返回一个token，用于区分不同请求的队列
    token = int(time.time() * 1000) % 100000
    return render_template("start.html", pic_dir=pic_dir, pic_cnt=pic_cnt, token=token)


# 前端页面加载完成时执行的函数，开始预测结果了
@app.route('/work_start')
def process_data():
    # 准备专属队列容器
    token = request.args.get("token")
    message_q_dict[token] = queue.Queue(maxsize=10)
    # 接受前台传过来的pic_dir值
    pic_dir = request.args.get('pic_dir')
    print("token=", token, pic_dir, "开始分析")
    # 开始预测
    index = 0
    for parent, _, files in os.walk(pic_dir):
        for file in files:
            if file[-3:] != 'jpg':
                continue
            # 预测瑕疵
            flaw_type, prob = predict(os.path.join(parent, file))
            # 将预测的结果存到消息队列里面
            path_after_dir = os.path.join(parent, file)[len(pic_dir):]
            message = {'file_path': path_after_dir, 'flaw_type': flaw_type, 'prob': prob, 'index': index + 1}
            print("token=", token, message)
            try:
                message_q_dict[token].put(message, timeout=3)
            except queue.Full:
                print("token=", token, "数据出队超时，页面可能中断了")
                return "ended"
            index += 1
    message_q_dict[token].put("ended")
    print("token=", token, "数据处理完毕")
    return "ended"


# 前端向后台请求一次结果时执行的代码
@app.route('/get_data')
def show_progress():
    token = request.args.get("token")
    try:
        # 从队列取一个消息并以json串的形式返回给前端
        q = message_q_dict[token].get_nowait()
    except queue.Empty:
        q = "queue is empty"
    return jsonify(q)


def predict(image_path):
    """
    输出图片，返回预测结果，label, pro
    """
    img = load_img(image_path, target_size=(in_size[0], in_size[1], 3), interpolation='antialias')
    img = img_to_array(img)
    img = preprocess_input(img)
    img = np.expand_dims(img, axis=0)
    res = model.predict(img)[0]
    # print(res)
    max_i = int(np.argmax(res))
    label = flaw_names[max_i]
    pro = "{:.2f}%".format(res[max_i] * 100)
    return label, pro


if __name__ == '__main__':
    # 规定开头，点击小箭头开始代码
    app.run()
